# import argparse
# import os
# import glob
# import cv2
# import numpy as np
# import torch
# from basicsr.archs.dualCodeBook_arch import DualCodeBookNet

# def load_model(model_path, device):
#     """根据参数加载DualCodeBookNet模型"""
#     model = DualCodeBookNet(
#         in_channel=3,
#         codebook_params=[[32, 1024, 512]],
#         gt_resolution=256,
#         LQ_stage=True,
#         use_quantize=True,
#         use_semantic_loss=False,
#         use_residual=False,
#         scale_factor=1,
#         former=False,
#         quantize_way='attention',
#     )
#     # 加载模型权重
#     ckpt = torch.load(model_path, map_location=device)
#     # 权重文件可能是一个字典，包含 'params' 键
#     if isinstance(ckpt, dict) and 'params' in ckpt:
#         state = ckpt['params']
#     else:
#         state = ckpt
#     model.load_state_dict(state, strict=False)
#     model.eval().to(device)
#     return model

# def to_tensor(img_bgr):
#     """[H,W,3] BGR float32 [0,1] -> torch [1,3,H,W] RGB"""
#     rgb = img_bgr[:, :, ::-1] # BGR to RGB
#     tensor = torch.from_numpy(np.ascontiguousarray(np.transpose(rgb, (2, 0, 1)))).float().unsqueeze(0)
#     return tensor


# def main():
#     parser = argparse.ArgumentParser()
#     parser.add_argument('--model_path', type=str,
#                         default='/root/code/fe-ma-sr/experiments/HQ_attention_archived_20251109_030552/models/net_g_best_.pth',
#                         help='预训练模型权重路径')
#     parser.add_argument('--hazy_root', type=str, default='/root/data/Outdoor/SOTS_restoration/SOTS/outdoor/hazy/', help='有雾图像输入文件夹')
#     parser.add_argument('--gt_root', type=str, default='/root/data/Outdoor/SOTS_restoration/SOTS/outdoor/gt/', help='清晰图像(GT)输入文件夹')
#     parser.add_argument('--output', type=str, default='/root/code/fe-ma-sr/result-dehaze-restored-dualbook-attention/image', help='结果输出文件夹')
#     args = parser.parse_args()

#     os.makedirs(args.output, exist_ok=True)
#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

#     # 加载模型
#     model = load_model(args.model_path, device)
#     print(f'模型已加载: {args.model_path}')

#     hazy_list = sorted(glob.glob(os.path.join(args.hazy_root, '*')))
#     if not hazy_list:
#         print(f'未在文件夹中找到任何有雾图像: {args.hazy_root}')
#         # return

#     print(f'开始处理 {len(hazy_list)} 张图像...')
#     codebook_usage = 0
#     processed_count = 0
#     for idx, hazy_path in enumerate(hazy_list):
#         hazy_name = os.path.basename(hazy_path)
        
#         # [逻辑修改] 假设hazy和gt文件夹下图片名称相同，直接构建gt路径
#         gt_path = os.path.join(args.gt_root, hazy_name)

#         # [逻辑修改] 检查对应的gt文件是否存在
#         if not os.path.exists(gt_path):
#             print(f'[{idx+1}/{len(hazy_list)}] 跳过 {hazy_name}: 在gt文件夹中未找到同名文件')
#             continue

#         print(f'[{idx+1}/{len(hazy_list)}] 正在处理: {hazy_name}')

#         img_hazy = cv2.imread(hazy_path, cv2.IMREAD_COLOR)
#         img_gt = cv2.imread(gt_path, cv2.IMREAD_COLOR)
#         if img_hazy is None or img_gt is None:
#             print(f'  -> 读取图像失败: {hazy_name} 或 {gt_path}')
#             continue
        
#         # 检查hazy和gt图像的尺寸是否一致
#         if img_hazy.shape[:2] != img_gt.shape[:2]:
#             print(f'  -> 跳过: 有雾图像和GT图像尺寸不匹配: {img_hazy.shape[:2]} vs {img_gt.shape[:2]}')
#             continue

#         # 归一化到 [0, 1]
#         img_hazy = img_hazy.astype(np.float32) / 255.0
#         img_gt = img_gt.astype(np.float32) / 255.0

#         # 将原始图像转换为Tensor
#         t_hazy = to_tensor(img_hazy).to(device)
#         t_gt = to_tensor(img_gt).to(device)

#         # 模型推理
#         with torch.no_grad():
#             out_hazy, out_clear, _, _, codebook_usage_tensor = model(t_gt, None)

#         # 处理并保存clear分支的输出结果
#         out_clear_np = out_clear.float().cpu().clamp_(0, 1).numpy()[0]
#         out_clear_np = np.transpose(out_clear_np, (1, 2, 0))[:, :, ::-1]  # RGB->BGR
#         out_clear_np = (out_clear_np * 255.0).round().astype(np.uint8)
#         save_name = os.path.splitext(hazy_name)[0] + '_restored.png'
#         cv2.imwrite(os.path.join(args.output, save_name), out_clear_np)
        
#         # 累加每次的codebook使用率
#         if codebook_usage_tensor is not None:
#             codebook_usage += codebook_usage_tensor.item()
#             processed_count += 1

#         if out_hazy is not None:
#             # 处理并保存 hazy 分支的输出结果
#             out_hazy_np = out_hazy.float().cpu().clamp_(0, 1).numpy()[0]
#             out_hazy_np = np.transpose(out_hazy_np, (1, 2, 0))[:, :, ::-1]  # RGB -> BGR
#             out_hazy_np = (out_hazy_np * 255.0).round().astype(np.uint8)
#             save_name_hazy = os.path.splitext(hazy_name)[0] + '_hazy_restored.png'
#             cv2.imwrite(os.path.join(args.output, save_name_hazy), out_hazy_np)

#     print(f'全部完成。结果已保存至: {args.output}')
#     if processed_count > 0:
#         codebook_usage_avg = codebook_usage / processed_count
#         print(f'处理了 {processed_count} 张图像，平均代码本使用率: {codebook_usage_avg:.4f}')
#     else:
#         print("没有处理任何图像。")

# if __name__ == '__main__':
#     main()

import argparse
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import glob
import cv2
import numpy as np
import torch
from basicsr.archs.dualCodeBook_arch import DualCodeBookNet

def load_model(model_path, device):
    """根据参数加载DualCodeBookNet模型"""
    model = DualCodeBookNet(
        in_channel=3,
        codebook_params=[[32, 1024, 512]],
        gt_resolution=256,
        LQ_stage=False,
        use_quantize=True,
        use_semantic_loss=False,
        use_residual=False,
        scale_factor=1,
        former=False,
        quantize_way='min',
        show_codebook_usage=False,
    )
    # 加载模型权重
    ckpt = torch.load(model_path, map_location=device)
    # 权重文件可能是一个字典，包含 'params' 键
    if isinstance(ckpt, dict) and 'params' in ckpt:
        state = ckpt['params']
    else:
        state = ckpt
    model.load_state_dict(state, strict=False)
    model.eval().to(device)
    return model

# === 代码修改 1: 更新 to_tensor 函数 ===
# 由于预处理步骤已经将图像转为RGB，此函数不再需要进行BGR->RGB的转换
def to_tensor(img_rgb):
    """[H,W,3] RGB float32 [0,1] -> torch [1,3,H,W] RGB"""
    # 输入的img已经是RGB格式，直接转置通道即可
    tensor = torch.from_numpy(np.ascontiguousarray(np.transpose(img_rgb, (2, 0, 1)))).float().unsqueeze(0)
    return tensor


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str,
                        default='/root/code/fe-ma-sr/experiments/HQ_Low_light_KL_0.05/models/net_g_best_.pth',
                        help='预训练模型权重路径')
    parser.add_argument('--hazy_root', type=str, default='/root/data/LOLdataset/eval15/low/', help='有雾图像输入文件夹')
    parser.add_argument('--gt_root', type=str, default='/root/data/LOLdataset/eval15/high/', help='清晰图像(GT)输入文件夹')
    parser.add_argument('--output', type=str, default='/root/code/fe-ma-sr/result-lowlight-restored-dualbook-min/kl_0.05_images', help='结果输出文件夹')
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载模型
    model = load_model(args.model_path, device)
    print(f'模型已加载: {args.model_path}')

    hazy_list = sorted(glob.glob(os.path.join(args.hazy_root, '*')))
    if not hazy_list:
        print(f'未在文件夹中找到任何有雾图像: {args.hazy_root}')
        return # 修改：如果列表为空，直接退出

    print(f'开始处理 {len(hazy_list)} 张图像...')
    codebook_usage = 0
    processed_count = 0
    for idx, hazy_path in enumerate(hazy_list):
        hazy_name = os.path.basename(hazy_path)
        
        # 假设hazy和gt文件夹下图片名称相同，直接构建gt路径
        gt_path = os.path.join(args.gt_root, hazy_name)

        # 检查对应的gt文件是否存在
        if not os.path.exists(gt_path):
            print(f'[{idx+1}/{len(hazy_list)}] 跳过 {hazy_name}: 在gt文件夹中未找到同名文件')
            continue

        print(f'[{idx+1}/{len(hazy_list)}] 正在处理: {hazy_name}')

        # === 代码修改 2: 应用指定的预处理方式 ===
        # 使用 try-except 块来处理 cv2.imread 返回 None 后调用 .astype 导致的错误
        try:
            # 1. 使用 OpenCV 读取 BGR 图像, 转换为 float32 并归一化
            img_hazy = cv2.imread(hazy_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
            img_gt = cv2.imread(gt_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.

            # 2. 将 BGR 图像转换为 RGB
            img_hazy = img_hazy[:, :, [2, 1, 0]]
            img_gt = img_gt[:, :, [2, 1, 0]]

        except AttributeError:
            # 如果 cv2.imread 读取失败返回 None, .astype 会触发 AttributeError
            print(f'  -> 读取图像失败或文件损坏: {hazy_name} 或 {gt_path}')
            continue
        
        # 检查hazy和gt图像的尺寸是否一致 (这一步在转换后进行)
        if img_hazy.shape[:2] != img_gt.shape[:2]:
            print(f'  -> 跳过: 有雾图像和GT图像尺寸不匹配: {img_hazy.shape[:2]} vs {img_gt.shape[:2]}')
            continue
        # === 预处理修改结束 ===

        # 将原始图像转换为Tensor (现在传入的是RGB图像)
        t_hazy = to_tensor(img_hazy).to(device)
        t_gt = to_tensor(img_gt).to(device)

        # 模型推理
        with torch.no_grad():
            # 注意：原始代码中模型输入为 (t_gt, None)，这里保持一致。
            # 如果模型应该输入 t_hazy，请修改为 model(t_hazy, t_gt) 或其他正确形式
            out_hazy, out_clear, _, _, codebook_usage_tensor, _ = model(t_hazy, t_gt) 

        # 处理并保存clear分支的输出结果
        out_clear_np = out_clear.float().cpu().clamp_(0, 1).numpy()[0]
        out_clear_np = np.transpose(out_clear_np, (1, 2, 0))[:, :, ::-1]  # Tensor(RGB) -> HWC(RGB) -> HWC(BGR)
        out_clear_np = (out_clear_np * 255.0).round().astype(np.uint8)
        save_name = os.path.splitext(hazy_name)[0] + '_restored.png'
        cv2.imwrite(os.path.join(args.output, save_name), out_clear_np)
        
        # 累加每次的codebook使用率
        if codebook_usage_tensor is not None:
            codebook_usage += codebook_usage_tensor.item()
            processed_count += 1

        if out_hazy is not None:
            # 处理并保存 hazy 分支的输出结果
            out_hazy_np = out_hazy.float().cpu().clamp_(0, 1).numpy()[0]
            out_hazy_np = np.transpose(out_hazy_np, (1, 2, 0))[:, :, ::-1]  # Tensor(RGB) -> HWC(RGB) -> HWC(BGR)
            out_hazy_np = (out_hazy_np * 255.0).round().astype(np.uint8)
            save_name_hazy = os.path.splitext(hazy_name)[0] + '_hazy_restored.png'
            cv2.imwrite(os.path.join(args.output, save_name_hazy), out_hazy_np)

    print(f'全部完成。结果已保存至: {args.output}')
    if processed_count > 0:
        codebook_usage_avg = codebook_usage / processed_count
        print(f'处理了 {processed_count} 张图像，平均代码本使用率: {codebook_usage_avg:.4f}')
    else:
        print("没有处理任何图像。")

if __name__ == '__main__':
    main()